#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@File    ：node.py
@Author  ：平
@Date    ：2025/9/28 17:08 
"""
import json
from typing import Literal

from langchain_core.messages import HumanMessage, AIMessage

from app.graph.agent import get_agent
from app.graph.llm import get_llm
from app.graph.prompt.prompt import get_prompt
from app.graph.state import State
from langgraph.types import Command
from datetime import datetime

from app.graph.tool import handler_to_retrieval, rag_retrieval, web_search, web_page, query_book
from app.graph.types import CompleteContent


def coordinator(state: State) -> Command[Literal['__end__', 'retrieval', 'complete']]:
    """
    coordinator协调智能体
    """
    if state.get("is_fill", False) and state.get("book", {}):
        return Command(goto='complete')

    # 创建协调者智能体
    coordinator_model = get_llm().bind_tools([handler_to_retrieval])
    # 获取提示词
    messages = get_prompt("coordinator", current_time=datetime.now())
    messages.extend(state.get("messages", []))
    goto = '__end__'
    intention = ""
    # 调用模型
    response = coordinator_model.invoke(messages)
    # 如果调用了handler_to_retrieval工具，将执行权交接给执行者
    if response.tool_calls:
        for tool_call in response.tool_calls:
            if tool_call.get("name", "") == "handler_to_retrieval":
                goto = 'retrieval'
                if tool_call.get("args", {}).get("intention", ""):
                    intention = tool_call.get("args", {}).get("intention", "")
                break

    messages = state.get("messages", [])
    if response.content:
        messages.append(AIMessage(content=response.content, name="coordinator"))

    return Command(
        update={
            "messages": messages,
            "intention": intention
        },
        goto=goto)


async def complete(state: State):
    """
    complete补全智能体
    """
    # 创建智能体
    agent = get_agent(name="complete", response_format=CompleteContent, tools=[web_search, web_page])
    # 获取系统提示词
    messages = get_prompt("complete", book=json.dumps(state.get("book", {}), ensure_ascii=False, indent=4))
    response = await agent.ainvoke({"messages": messages})
    structured_resp = response.get("structured_response", {})
    if hasattr(structured_resp, "model_dump"):
        structured_resp = structured_resp.model_dump()
    elif hasattr(structured_resp, '__dict__'):
        structured_resp = structured_resp.__dict__
    else:
        structured_resp = {}
    book = {
        **state.get("book", {}),
        **structured_resp
    }
    return {
        "messages": response.get("messages", []),
        "book": book
    }


async def retrieval(state: State):
    """
    retrieval检索智能体
    """
    pass
    # 创建智能体
    agent = get_agent(name="retrieval", tools=[rag_retrieval, query_book, web_search, web_page])
    # 获取提示词
    messages = get_prompt("retrieval", intention=state.get("intention", ""))
    messages.extend(state.get("messages", []))
    response = await agent.ainvoke({"messages": messages})
    return {
        "messages": response.get("messages", [])
    }


if __name__ == '__main__':
    pass

    # llm = DEFAULT_LLM.bind_tools([handler_to_retrieval])
    # messages = get_prompt("coordinator", current_time=datetime.now())
    # messages.append(HumanMessage(content="你好，当前时间"))
    # response = llm.invoke(messages)
    # print(response)
    # async def main():
    #     # 创建智能体
    #     agent = get_agent(name="complete", response_format=CompleteContent, tools=[web_search, web_page])
    #     # 获取系统提示词
    #     messages = get_prompt("complete", book=json.dumps({"id": 123, "title": "《三体》"}, ensure_ascii=False, indent=4))
    #     response = await agent.ainvoke({"messages": messages})
    #     print(response)
    #
    #
    # asyncio.run(main())

    # async def main():
    #     # 创建智能体
    #     agent = get_agent(name="retrieval", tools=[rag_retrieval, web_search, web_page])
    #     # 获取提示词
    #     messages = get_prompt("retrieval", intention="获取《三体》书籍在讲什么")
    #     messages.append(HumanMessage(content="《三体》书籍在讲什么"))
    #     response = await agent.ainvoke({"messages": messages})
    #     print(response)
    #
    # asyncio.run(main())
